import numpy as np
import matplotlib.pyplot as plt
import scipy.io as sio
import os, sys

def load_target_mat(target_file):
    data_dict = sio.loadmat(target_file)

    lh_data = data_dict['lh'].transpose()[0]
    rh_data = data_dict['rh'].transpose()[0]

    return lh_data, rh_data

def cal_corr(src_list, des_list):
    corr = np.corrcoef(np.array(src_list), np.array(des_list))
    return corr[0][1]

def plot_hist(output_dir, ori_data, des_data, ori_label, des_label, symbol):
    fig,ax=plt.subplots()

    ax.hist(ori_data, bins=16, histtype="bar", alpha=0.6, label=ori_label)
    ax.hist(des_data, bins=16, histtype="bar", alpha=0.6, label=des_label)
    ax.legend()

    ax.set_xticks(np.arange(0.5, 1.0, 0.05))
    ax.set_xlim(0.5, 1.0)
    ax.set_yticks(np.arange(0, 1100, 100))

    ax.set_xlabel('Variability')
    ax.set_ylabel('Counts')
    ax.set_title('Site ID: ' + symbol)

    fig.tight_layout()

    plt.savefig(output_dir + symbol + '_' + ori_label + '_'+ des_label + '.png', dpi=150)


def plot_simple_hist(output_dir, ori_data, ori_label, symbol):
    fig,ax=plt.subplots()

    ax.hist(ori_data, bins=16, histtype="bar", alpha=0.6, label=ori_label)
    ax.legend()

    ax.set_xticks(np.arange(0.5, 1.0, 0.05))
    ax.set_xlim(0.5, 1.0)
    ax.set_yticks(np.arange(0, 1100, 100))

    ax.set_xlabel('Variability')
    ax.set_ylabel('Counts')
    ax.set_title('Site ID: ' + symbol)

    fig.tight_layout()

    plt.savefig(output_dir + symbol + '_' + ori_label + '.png', dpi=150)

def draw_one_fig_asd_hc(output_dir, variability_dir, asd_file, hc_file, symbol):
    asd_file = os.path.join(variability_dir, asd_file)
    hc_file = os.path.join(variability_dir, hc_file)

    lh_asd_data, rh_asd_data = load_target_mat(asd_file)
    lh_hc_data, rh_hc_data = load_target_mat(hc_file)

    wh_asd_data = []
    wh_asd_data.extend(lh_asd_data)
    wh_asd_data.extend(rh_asd_data)

    wh_hc_data = []
    wh_hc_data.extend(lh_hc_data)
    wh_hc_data.extend(rh_hc_data)

    plot_hist(output_dir, wh_asd_data, wh_hc_data, 'ASD', 'HC', symbol)

def draw_one_fig(output_dir, variability_file, symbol):
    lh_asd_data, rh_asd_data = load_target_mat(variability_file)

    wh_asd_data = []
    wh_asd_data.extend(lh_asd_data)
    wh_asd_data.extend(rh_asd_data)

    plot_simple_hist(output_dir, wh_asd_data, 'HC', symbol)

def fetch_data_label_from_file(src_file):
    lh_data, rh_data = load_target_mat(src_file)
    res_data = []
    res_data.extend(lh_data)
    res_data.extend(rh_data)

    file_name = src_file.split('.')[0]
    data_type = file_name.split('_')[-2]
    center_label = file_name.split('_')[-1]

    return res_data, center_label, data_type

def draw_tri_diff_fig(output_dir, fir_file, sec_file, tri_file, for_file, symbol):
    fir_data, fir_label, _ = fetch_data_label_from_file(fir_file)
    sec_data, sec_label, _ = fetch_data_label_from_file(sec_file)
    tri_data, tri_label, _ = fetch_data_label_from_file(tri_file)
    for_data, for_label, _ = fetch_data_label_from_file(for_file)

    # symbol = "_".join([fir_label, sec_label, tri_label, for_label])

    plot_tri_hist(output_dir, fir_data, sec_data, tri_data, for_data, fir_label, sec_label, tri_label, for_label, symbol)

def draw_multi_diff_fig(output_dir, file_list, symbol):
    fir_data, fir_label, _ = fetch_data_label_from_file(fir_file)
    sec_data, sec_label, _ = fetch_data_label_from_file(sec_file)
    tri_data, tri_label, _ = fetch_data_label_from_file(tri_file)
    for_data, for_label, _ = fetch_data_label_from_file(for_file)

    # symbol = "_".join([fir_label, sec_label, tri_label, for_label])
    plot_tri_hist(output_dir, fir_data, sec_data, tri_data, for_data, fir_label, sec_label, tri_label, for_label, symbol)

def plot_tri_hist(output_dir, ori_data, des_data, tri_data, for_data, ori_label, des_label, tri_label, for_label, symbol):
    fig,ax=plt.subplots()

    ax.hist(ori_data, bins=16, histtype="bar", alpha=0.75, label=ori_label)
    ax.hist(des_data, bins=16, histtype="bar", alpha=0.75, label=des_label)
    ax.hist(tri_data, bins=16, histtype="bar", alpha=0.75, label=tri_label)
    ax.hist(for_data, bins=16, histtype="bar", alpha=0.75, label=for_label)
    ax.legend()

    ax.set_xticks(np.arange(0.5, 1.0, 0.05))
    ax.set_xlim(0.5, 1.0)
    ax.set_yticks(np.arange(0, 1100, 100))

    ax.set_xlabel('Variability')
    ax.set_ylabel('Counts')
    ax.set_title('Site ID: ' + symbol)

    fig.tight_layout()

    oupput_path = os.path.join(output_dir, symbol + '.png')
    plt.savefig(oupput_path, dpi=150)

project_dir = '/mri_projects/ASD_Analysis'
variability_dir = os.path.join(project_dir, 'Results/Variability/')
output_dir = os.path.join(project_dir, 'Figures/Variability/')

# draw_one_fig_asd_hc(output_dir, variability_dir, "GU")
# draw_one_fig_asd_hc(output_dir, variability_dir, "KKI")
gu_hc_file = os.path.join(variability_dir, 'intra_Variability_FS4_HC_GU.mat')
kki_hc_file = os.path.join(variability_dir, 'intra_Variability_FS4_HC_KKI.mat')
ohsu_hc_file = os.path.join(variability_dir, 'inter_Variability_FS4_HC_OHSU.mat')
std_var_file = '/home/li/Documents/Trainning/SADTestFCMap/Scripts/Variability_FS4.mat'
# draw_one_fig(output_dir, std_var_file, '23-subjs')
# (output_dir, variability_dir, asd_file, hc_file, symbol)
ohsu_asd_file = os.path.join(variability_dir, 'inter_Variability_FS4_ASD_OHSU.mat')
ohsu_hc_file = os.path.join(variability_dir, 'inter_Variability_FS4_HC_OHSU.mat')
# draw_one_fig_asd_hc(output_dir, variability_dir, ohsu_asd_file, ohsu_hc_file, "OHSU")
# draw_tri_diff_fig(output_dir, gu_hc_file, kki_hc_file, ohsu_hc_file, std_var_file, 'GU_KKI_OHSU_23subjs')

gu_asd_file = os.path.join(variability_dir, 'inter_Variability_FS4_ASD_GU.mat')
kki_asd_file = os.path.join(variability_dir, 'inter_Variability_FS4_ASD_KKI.mat')
ohsu_asd_file = os.path.join(variability_dir, 'inter_Variability_FS4_ASD_OHSU.mat')
nyu_asd_file = os.path.join(variability_dir, 'inter_Variability_FS4_ASD_NYU.mat')
draw_tri_diff_fig(output_dir, gu_asd_file, kki_asd_file, ohsu_asd_file, nyu_asd_file, 'GU_KKI_OHSU_NYU_ASD')
